// Copyright lowRISC contributors (OpenTitan project).
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

#include "sw/device/tests/penetrationtests/firmware/lib/pentest_lib.h"

#include "hw/ip/aes/model/aes.h"
#include "sw/device/lib/arch/device.h"
#include "sw/device/lib/base/bitfield.h"
#include "sw/device/lib/base/csr.h"
#include "sw/device/lib/base/macros.h"
#include "sw/device/lib/base/mmio.h"
#include "sw/device/lib/base/status.h"
#include "sw/device/lib/dif/dif_alert_handler.h"
#include "sw/device/lib/dif/dif_clkmgr.h"
#include "sw/device/lib/dif/dif_csrng.h"
#include "sw/device/lib/dif/dif_csrng_shared.h"
#include "sw/device/lib/dif/dif_gpio.h"
#include "sw/device/lib/dif/dif_lc_ctrl.h"
#include "sw/device/lib/dif/dif_pinmux.h"
#include "sw/device/lib/dif/dif_rstmgr.h"
#include "sw/device/lib/dif/dif_rv_core_ibex.h"
#include "sw/device/lib/dif/dif_rv_plic.h"
#include "sw/device/lib/dif/dif_rv_timer.h"
#include "sw/device/lib/dif/dif_uart.h"
#include "sw/device/lib/runtime/hart.h"
#include "sw/device/lib/runtime/irq.h"
#include "sw/device/lib/runtime/print.h"
#include "sw/device/lib/testing/pinmux_testutils.h"
#include "sw/device/lib/testing/rv_plic_testutils.h"
#include "sw/device/lib/testing/test_framework/check.h"
#include "sw/device/lib/testing/test_framework/ottf_isrs.h"
#include "sw/device/lib/testing/test_framework/ottf_main.h"

#include "clkmgr_regs.h"  // Generated
#include "csrng_regs.h"   // Generated
#include "hw/top_earlgrey/sw/autogen/top_earlgrey.h"

#if !OT_IS_ENGLISH_BREAKFAST
#include "sw/device/lib/crypto/drivers/otbn.h"
#include "sw/device/lib/dif/dif_csrng.h"
#include "sw/device/lib/dif/dif_edn.h"
#include "sw/device/lib/dif/dif_entropy_src.h"
#include "sw/device/lib/testing/alert_handler_testutils.h"
#include "sw/device/lib/testing/entropy_testutils.h"

#include "edn_regs.h"  // Generated
#endif

/**
 * Bitfield for the trigger source.
 *
 * Bits 10 and 11 are used to select the trigger source. See chiplevel.sv.tpl
 * for details.
 */
static const bitfield_field32_t kTriggerSourceBitfield = {
    .index = 10,
    .mask = 0x3,
};

enum {
  /**
   * Bit index of the hardware trigger gate signal for gating the hardware
   * trigger from software.
   *
   * See chiplevel.sv.tpl for details.
   */
  kTriggerHwGateBitIndex = 9,
  /**
   * Bit index of the software trigger signal.
   *
   * See chiplevel.sv.tpl for details.
   */
  kTriggerSwBitIndex = 8,
  /**
   * RV timer settings.
   */
  kRvTimerComparator = 0,
  kRvTimerHart = kTopEarlgreyPlicTargetIbex0,
};

// By default, we use the precise, hardware-gated capture trigger.
static unsigned int trigger_bit_index = kTriggerHwGateBitIndex;

static dif_gpio_t gpio;
static dif_pinmux_t pinmux;
static dif_rv_timer_t timer;
static dif_uart_t uart0;
static dif_uart_t uart1;

#if !OT_IS_ENGLISH_BREAKFAST
static dif_alert_handler_t alert_handler;
static dif_csrng_t csrng;
static dif_edn_t edn0;
static dif_edn_t edn1;
static dif_lc_ctrl_t lc;
static dif_rv_plic_t plic;
static dif_rstmgr_t rstmgr;
#endif

status_t pentest_configure_entropy_source_max_reseed_interval(void) {
#if !OT_IS_ENGLISH_BREAKFAST
  const dif_entropy_src_t entropy_src = {
      .base_addr = mmio_region_from_addr(TOP_EARLGREY_ENTROPY_SRC_BASE_ADDR)};
  const dif_csrng_t csrng = {
      .base_addr = mmio_region_from_addr(TOP_EARLGREY_CSRNG_BASE_ADDR)};
  const dif_edn_t edn0 = {
      .base_addr = mmio_region_from_addr(TOP_EARLGREY_EDN0_BASE_ADDR)};
  const dif_edn_t edn1 = {
      .base_addr = mmio_region_from_addr(TOP_EARLGREY_EDN1_BASE_ADDR)};

  TRY(entropy_testutils_stop_all());

  // Re-eanble entropy src and csrng.
  TRY(dif_entropy_src_configure(
      &entropy_src, entropy_testutils_config_default(), kDifToggleEnabled));
  TRY(dif_csrng_configure(&csrng));

  // Re-enable EDN0 in auto mode.
  TRY(dif_edn_set_auto_mode(
      &edn0,
      (dif_edn_auto_params_t){
          // EDN0 provides lower-quality entropy.  Let one generate command
          // return 8
          // blocks, and reseed every 32 generates.
          .instantiate_cmd =
              {
                  .cmd = csrng_cmd_header_build(kCsrngAppCmdInstantiate,
                                                kDifCsrngEntropySrcToggleEnable,
                                                /*cmd_len=*/0,
                                                /*generate_len=*/0),
                  .seed_material =
                      {
                          .len = 0,
                      },
              },
          .reseed_cmd =
              {
                  .cmd = csrng_cmd_header_build(
                      kCsrngAppCmdReseed, kDifCsrngEntropySrcToggleEnable,
                      /*cmd_len=*/0, /*generate_len=*/0),
                  .seed_material =
                      {
                          .len = 0,
                      },
              },
          .generate_cmd =
              {
                  // Generate 8 128-bit blocks.
                  .cmd = csrng_cmd_header_build(kCsrngAppCmdGenerate,
                                                kDifCsrngEntropySrcToggleEnable,
                                                /*cmd_len=*/0,
                                                /*generate_len=*/8),
                  .seed_material =
                      {
                          .len = 0,
                      },
              },
          // Reseed every 0xffffffff generates.
          .reseed_interval = 0xffffffff,
      }));

  // Re-enable EDN1 in auto mode.
  TRY(dif_edn_set_auto_mode(
      &edn1,
      (dif_edn_auto_params_t){
          // EDN1 provides highest-quality entropy.  Let one generate command
          // return 1 block, and reseed after every generate.
          .instantiate_cmd =
              {
                  .cmd = csrng_cmd_header_build(kCsrngAppCmdInstantiate,
                                                kDifCsrngEntropySrcToggleEnable,
                                                /*cmd_len=*/0,
                                                /*generate_len=*/0),
                  .seed_material =
                      {
                          .len = 0,
                      },
              },
          .reseed_cmd =
              {
                  .cmd = csrng_cmd_header_build(
                      kCsrngAppCmdReseed, kDifCsrngEntropySrcToggleEnable,
                      /*cmd_len=*/0, /*generate_len=*/0),
                  .seed_material =
                      {
                          .len = 0,
                      },
              },
          .generate_cmd =
              {
                  // Generate 1 128-bit block.
                  .cmd = csrng_cmd_header_build(kCsrngAppCmdGenerate,
                                                kDifCsrngEntropySrcToggleEnable,
                                                /*cmd_len=*/0,
                                                /*generate_len=*/1),
                  .seed_material =
                      {
                          .len = 0,
                      },
              },
          // Reseed after every 0xffffffff generates.
          .reseed_interval = 0xffffffff,
      }));
#endif
  return OK_STATUS();
}

pentest_registered_alerts_t pentest_get_triggered_alerts(void) {
  pentest_registered_alerts_t registered;
  memset(registered.alerts, 0, sizeof(registered.alerts));

#if !OT_IS_ENGLISH_BREAKFAST
  bool is_cause;
  // Loop over all alert_cause regs
  for (size_t alert = 0; alert < ALERT_HANDLER_PARAM_N_ALERTS; alert++) {
    CHECK_DIF_OK(
        dif_alert_handler_alert_is_cause(&alert_handler, alert, &is_cause));
    if (is_cause) {
      if (alert < 32) {
        registered.alerts[0] |= (1 << alert);
      } else if (alert < 64) {
        registered.alerts[1] |= (1 << (alert - 32));
      } else {
        registered.alerts[2] |= (1 << (alert - 64));
      }
    }
  }

  // Loop over all alert_cause regs.
  for (dif_alert_handler_alert_t i = 0; i < ALERT_HANDLER_PARAM_N_ALERTS; i++) {
    CHECK_DIF_OK(dif_alert_handler_alert_acknowledge(&alert_handler, i));
  }
#endif

  return registered;
}

void pentest_configure_alert_handler(void) {
#if !OT_IS_ENGLISH_BREAKFAST
  irq_global_ctrl(true);
  irq_external_ctrl(true);

  mmio_region_t base_addr =
      mmio_region_from_addr(TOP_EARLGREY_RV_PLIC_BASE_ADDR);
  CHECK_DIF_OK(dif_rv_plic_init(base_addr, &plic));

  base_addr = mmio_region_from_addr(TOP_EARLGREY_ALERT_HANDLER_BASE_ADDR);
  CHECK_DIF_OK(dif_alert_handler_init(base_addr, &alert_handler));

  CHECK_DIF_OK(dif_rstmgr_init(
      mmio_region_from_addr(TOP_EARLGREY_RSTMGR_AON_BASE_ADDR), &rstmgr));

  dif_alert_handler_alert_t alerts[ALERT_HANDLER_PARAM_N_ALERTS];
  dif_alert_handler_class_t alert_classes[ALERT_HANDLER_PARAM_N_ALERTS];

  // Enable all incoming alerts and configure them to classa.
  for (dif_alert_handler_alert_t i = 0; i < ALERT_HANDLER_PARAM_N_ALERTS; ++i) {
    alerts[i] = i;
    alert_classes[i] = kDifAlertHandlerClassA;
  }

  dif_alert_handler_escalation_phase_t esc_phases[] = {
      {.phase = kDifAlertHandlerClassStatePhase0,
       .signal = 0,
       .duration_cycles = 2000}};

  dif_alert_handler_class_config_t class_config = {
      .auto_lock_accumulation_counter = kDifToggleDisabled,
      .accumulator_threshold = 0,
      .irq_deadline_cycles = 10000,
      .escalation_phases = esc_phases,
      .escalation_phases_len = ARRAYSIZE(esc_phases),
      .crashdump_escalation_phase = kDifAlertHandlerClassStatePhase1,
  };

  dif_alert_handler_class_config_t class_configs[] = {class_config};

  dif_alert_handler_class_t classes[] = {kDifAlertHandlerClassA};
  dif_alert_handler_config_t config = {
      .alerts = alerts,
      .alert_classes = alert_classes,
      .alerts_len = ARRAYSIZE(alerts),
      .classes = classes,
      .class_configs = class_configs,
      .classes_len = ARRAYSIZE(class_configs),
      .ping_timeout = 256,
  };

  CHECK_STATUS_OK(alert_handler_testutils_configure_all(&alert_handler, config,
                                                        kDifToggleEnabled));
  // Enables alert handler irq.
  CHECK_DIF_OK(dif_alert_handler_irq_set_enabled(
      &alert_handler, kDifAlertHandlerIrqClassa, kDifToggleEnabled));
#endif
}

status_t pentest_read_device_id(uint32_t device_id[]) {
#if !OT_IS_ENGLISH_BREAKFAST
  mmio_region_t lc_reg =
      mmio_region_from_addr(TOP_EARLGREY_LC_CTRL_REGS_BASE_ADDR);
  CHECK_DIF_OK(dif_lc_ctrl_init(lc_reg, &lc));

  dif_lc_ctrl_device_id_t lc_device_id;
  CHECK_DIF_OK(dif_lc_ctrl_get_device_id(&lc, &lc_device_id));
  memcpy(device_id, lc_device_id.data, 8 * sizeof(uint32_t));
#else
  memset(device_id, 0, 8 * sizeof(uint32_t));
#endif

  return OK_STATUS();
}

void pentest_configure_cpu(void) {
  uint32_t cpuctrl_csr;
  // Get current config.
  CSR_READ(CSR_REG_CPUCTRL, &cpuctrl_csr);
  // Disable the iCache.
  cpuctrl_csr = bitfield_field32_write(
      cpuctrl_csr, (bitfield_field32_t){.mask = 0x1, .index = 0}, 0);
  // Disable dummy instructions.
  cpuctrl_csr = bitfield_field32_write(
      cpuctrl_csr, (bitfield_field32_t){.mask = 0x1, .index = 2}, 0);
  // Write back config.
  CSR_WRITE(CSR_REG_CPUCTRL, cpuctrl_csr);
}

/**
 * Initializes the UART peripheral.
 */
static void pentest_init_uart(void) {
  CHECK(kUartBaudrate <= UINT32_MAX, "kUartBaudrate must fit in uint32_t");
  CHECK(kClockFreqPeripheralHz <= UINT32_MAX,
        "kClockFreqPeripheralHz must fit in uint32_t");
  const dif_uart_config_t uart_config = {
      .baudrate = (uint32_t)kUartBaudrate,
      .clk_freq_hz = (uint32_t)kClockFreqPeripheralHz,
      .parity_enable = kDifToggleDisabled,
      .parity = kDifUartParityEven,
      .tx_enable = kDifToggleEnabled,
      .rx_enable = kDifToggleEnabled,
  };

  OT_DISCARD(dif_uart_init(mmio_region_from_addr(TOP_EARLGREY_UART0_BASE_ADDR),
                           &uart0));
  OT_DISCARD(dif_uart_configure(&uart0, uart_config));
  base_uart_stdout(&uart0);

#if !OT_IS_ENGLISH_BREAKFAST
  OT_DISCARD(dif_uart_init(mmio_region_from_addr(TOP_EARLGREY_UART1_BASE_ADDR),
                           &uart1));
  OT_DISCARD(dif_uart_configure(&uart1, uart_config));
#endif
}

/**
 * Initializes the GPIO peripheral.
 *
 * @param trigger Trigger source.
 */
static void pentest_init_gpio(pentest_trigger_source_t trigger) {
  OT_DISCARD(
      dif_gpio_init(mmio_region_from_addr(TOP_EARLGREY_GPIO_BASE_ADDR), &gpio));

  uint32_t select_mask =
      bitfield_field32_write(0, kTriggerSourceBitfield, UINT32_MAX);
  uint32_t enable_mask = bitfield_bit32_write(0, kTriggerHwGateBitIndex, true);
  enable_mask = bitfield_bit32_write(enable_mask, kTriggerSwBitIndex, true);

  // Configure the pinmux to enable the GPIOs.
  for (size_t i = 0; i < 32; ++i) {
    if ((select_mask | enable_mask) & (1u << i)) {
      dif_pinmux_index_t mio = kTopEarlgreyPinmuxInselIoa0 + i;
      dif_pinmux_index_t periph_io =
          kTopEarlgreyPinmuxPeripheralInGpioGpio0 + i;
      OT_DISCARD(dif_pinmux_input_select(&pinmux, periph_io, mio));

      mio = kTopEarlgreyPinmuxMioOutIoa0 + i;
      periph_io = kTopEarlgreyPinmuxOutselGpioGpio0 + i;
      OT_DISCARD(dif_pinmux_output_select(&pinmux, mio, periph_io));
    }
  }

  OT_DISCARD(dif_gpio_output_set_enabled_all(&gpio, select_mask | enable_mask));

  OT_DISCARD(dif_gpio_write_masked(
      &gpio, select_mask,
      bitfield_field32_write(0, kTriggerSourceBitfield, trigger)));
}

/**
 * Initializes the timer peripheral.
 */
static void pentest_init_timer(void) {
  OT_DISCARD(dif_rv_timer_init(
      mmio_region_from_addr(TOP_EARLGREY_RV_TIMER_BASE_ADDR), &timer));
  OT_DISCARD(dif_rv_timer_reset(&timer));
  dif_rv_timer_tick_params_t tick_params;
  OT_DISCARD(dif_rv_timer_approximate_tick_params(
      kClockFreqPeripheralHz, kClockFreqCpuHz, &tick_params));
  OT_DISCARD(dif_rv_timer_set_tick_params(&timer, kRvTimerHart, tick_params));
  OT_DISCARD(dif_rv_timer_irq_set_enabled(
      &timer, kDifRvTimerIrqTimerExpiredHart0Timer0, kDifToggleEnabled));
  irq_timer_ctrl(true);
  irq_global_ctrl(true);
}

/**
 * Initializes the CSRNG handle.
 */
static void pentest_init_csrng(void) {
#if !OT_IS_ENGLISH_BREAKFAST
  OT_DISCARD(dif_csrng_init(mmio_region_from_addr(TOP_EARLGREY_CSRNG_BASE_ADDR),
                            &csrng));
  OT_DISCARD(dif_csrng_init(mmio_region_from_addr(TOP_EARLGREY_CSRNG_BASE_ADDR),
                            &csrng));
#endif
}

/**
 * Initializes the EDN handle.
 */
static void pentest_init_edn(void) {
#if !OT_IS_ENGLISH_BREAKFAST
  OT_DISCARD(
      dif_edn_init(mmio_region_from_addr(TOP_EARLGREY_EDN0_BASE_ADDR), &edn0));

  OT_DISCARD(
      dif_edn_init(mmio_region_from_addr(TOP_EARLGREY_EDN1_BASE_ADDR), &edn1));
#endif
}

/**
 * Override default Timer ISR.
 *
 * Disables the counter and clears pending interrupts.
 */
void ottf_timer_isr(uint32_t *exc_info) {
  // Return values of below functions are ignored to improve capture
  // performance.
  OT_DISCARD(dif_rv_timer_counter_set_enabled(&timer, kRvTimerHart,
                                              kDifToggleDisabled));
  OT_DISCARD(dif_rv_timer_irq_acknowledge(
      &timer, kDifRvTimerIrqTimerExpiredHart0Timer0));
}

/**
 * Disables the given peripherals to reduce noise during SCA.
 *
 * Care must be taken when disabling the entropy complex if a peripheral that
 * depends on it will be used in SCA. E.g., We can disable the entropy complex
 * when analyzing AES only because AES features a parameter to skip PRNG
 * reseeding for SCA experiments. Without this parameter, AES would simply get
 * stalled with a disabled entropy complex.
 *
 * @param disable Set of peripherals to disable.
 */
void sca_disable_peripherals(pentest_peripherals_t disable) {
#if !OT_IS_ENGLISH_BREAKFAST
  if (disable & kPentestPeripheralEdn) {
    OT_DISCARD(dif_edn_stop(&edn0));
    OT_DISCARD(dif_edn_stop(&edn1));
  }
  if (disable & kPentestPeripheralCsrng) {
    OT_DISCARD(dif_csrng_stop(&csrng));
  }
  if (disable & kPentestPeripheralEntropy) {
    dif_entropy_src_t entropy;
    OT_DISCARD(dif_entropy_src_init(
        mmio_region_from_addr(TOP_EARLGREY_ENTROPY_SRC_BASE_ADDR), &entropy));
    OT_DISCARD(dif_entropy_src_set_enabled(&entropy, kDifToggleDisabled));
  }
#endif

  // Disable HMAC, KMAC, OTBN and USB clocks through CLKMGR DIF.
  dif_clkmgr_t clkmgr;
  OT_DISCARD(dif_clkmgr_init(
      mmio_region_from_addr(TOP_EARLGREY_CLKMGR_AON_BASE_ADDR), &clkmgr));

  if (disable & kPentestPeripheralAes) {
    OT_DISCARD(dif_clkmgr_hintable_clock_set_hint(
        &clkmgr, CLKMGR_CLK_HINTS_CLK_MAIN_AES_HINT_BIT, kDifToggleDisabled));
  }
#if !OT_IS_ENGLISH_BREAKFAST
  if (disable & kPentestPeripheralHmac) {
    OT_DISCARD(dif_clkmgr_hintable_clock_set_hint(
        &clkmgr, CLKMGR_CLK_HINTS_CLK_MAIN_HMAC_HINT_BIT, kDifToggleDisabled));
  }
#endif
  if (disable & kPentestPeripheralIoDiv4) {
    OT_DISCARD(dif_clkmgr_gateable_clock_set_enabled(
        &clkmgr, CLKMGR_CLK_ENABLES_CLK_IO_DIV4_PERI_EN_BIT,
        kDifToggleDisabled));
  }
  if (disable & kPentestPeripheralIoDiv2) {
    OT_DISCARD(dif_clkmgr_gateable_clock_set_enabled(
        &clkmgr, CLKMGR_CLK_ENABLES_CLK_IO_DIV2_PERI_EN_BIT,
        kDifToggleDisabled));
  }
  if (disable & kPentestPeripheralUsb) {
    OT_DISCARD(dif_clkmgr_gateable_clock_set_enabled(
        &clkmgr, CLKMGR_CLK_ENABLES_CLK_USB_PERI_EN_BIT, kDifToggleDisabled));
  }
  if (disable & kPentestPeripheralIo) {
    OT_DISCARD(dif_clkmgr_gateable_clock_set_enabled(
        &clkmgr, CLKMGR_CLK_ENABLES_CLK_IO_PERI_EN_BIT, kDifToggleDisabled));
  }

#if !OT_IS_ENGLISH_BREAKFAST
  if (disable & kPentestPeripheralKmac) {
    OT_DISCARD(dif_clkmgr_hintable_clock_set_hint(
        &clkmgr, CLKMGR_CLK_HINTS_CLK_MAIN_KMAC_HINT_BIT, kDifToggleDisabled));
  }
  if (disable & kPentestPeripheralOtbn) {
    OT_DISCARD(dif_clkmgr_hintable_clock_set_hint(
        &clkmgr, CLKMGR_CLK_HINTS_CLK_MAIN_OTBN_HINT_BIT, kDifToggleDisabled));
  }
#endif
}

void pentest_init(pentest_trigger_source_t trigger,
                  pentest_peripherals_t enable) {
  OT_DISCARD(dif_pinmux_init(
      mmio_region_from_addr(TOP_EARLGREY_PINMUX_AON_BASE_ADDR), &pinmux));
  pinmux_testutils_init(&pinmux);
  pentest_init_uart();
  pentest_init_gpio(trigger);
  pentest_init_timer();
  pentest_init_csrng();
  pentest_init_edn();
  sca_disable_peripherals(~enable);
}

const dif_uart_t *pentest_get_uart(void) { return &uart1; }

void pentest_select_trigger_type(pentest_trigger_type_t trigger_type) {
  if (trigger_type == kPentestTriggerTypeHwGated) {
    trigger_bit_index = kTriggerHwGateBitIndex;
  } else if (trigger_type == kPentestTriggerTypeSw) {
    trigger_bit_index = kTriggerSwBitIndex;
  }
}

void pentest_set_trigger_high(void) {
  OT_DISCARD(dif_gpio_write(&gpio, trigger_bit_index, true));
}

void pentest_set_trigger_low(void) {
  OT_DISCARD(dif_gpio_write(&gpio, trigger_bit_index, false));
}

void pentest_call_and_sleep(sca_callee callee, uint32_t sleep_cycles,
                            bool sw_trigger, bool otbn) {
  // Disable the IO_DIV4_PERI clock to reduce noise during the actual capture.
  // This also disables the UART(s) and GPIO modules required for
  // communication with the scope. Therefore, it has to be re-enabled after
  // the capture.
  dif_clkmgr_t clkmgr;
  OT_DISCARD(dif_clkmgr_init(
      mmio_region_from_addr(TOP_EARLGREY_CLKMGR_AON_BASE_ADDR), &clkmgr));

  // Start timer to wake Ibex after the callee is done.
  uint64_t current_time;
  // Return values of below functions are ignored to improve capture
  // performance.
  OT_DISCARD(dif_rv_timer_counter_read(&timer, kRvTimerHart, &current_time));
  OT_DISCARD(dif_rv_timer_arm(&timer, kRvTimerHart, kRvTimerComparator,
                              current_time + sleep_cycles));
  OT_DISCARD(dif_rv_timer_counter_set_enabled(&timer, kRvTimerHart,
                                              kDifToggleEnabled));

  if (sw_trigger) {
    pentest_set_trigger_high();
  }

  callee();

  wait_for_interrupt();

  // #if !OT_IS_ENGLISH_BREAKFAST
  //   if (otbn) {
  //     otbn_busy_wait_for_done();
  //   }
  // #endif

  if (sw_trigger) {
    pentest_set_trigger_low();
  }

  // Re-enable IO_DIV4_PERI clock to resume communication with the scope.
  OT_DISCARD(dif_clkmgr_gateable_clock_set_enabled(
      &clkmgr, CLKMGR_CLK_ENABLES_CLK_IO_DIV4_PERI_EN_BIT, kDifToggleEnabled));
}

static uint32_t sca_lfsr_state_masking = 0xdeadbeef;
static uint32_t sca_lfsr_state_order = 0x99999999;

void pentest_seed_lfsr(uint32_t seed, pentest_lfsr_context_t context) {
  if (context == kPentestLfsrMasking) {
    sca_lfsr_state_masking = seed;
  }
  if (context == kPentestLfsrOrder) {
    sca_lfsr_state_order = seed;
  }
}

uint32_t pentest_next_lfsr(uint16_t num_steps, pentest_lfsr_context_t context) {
  uint32_t sca_lfsr_state;
  if (context == kPentestLfsrMasking) {
    sca_lfsr_state = sca_lfsr_state_masking;
  }
  if (context == kPentestLfsrOrder) {
    sca_lfsr_state = sca_lfsr_state_order;
  }
  const uint32_t lfsr_out = sca_lfsr_state;
  for (size_t i = 0; i < num_steps; ++i) {
    bool lfsr_bit = sca_lfsr_state & 0x00000001;
    sca_lfsr_state = sca_lfsr_state >> 1;
    if (lfsr_bit) {
      sca_lfsr_state ^= 0x80000057;
    }
  }
  if (context == kPentestLfsrMasking) {
    sca_lfsr_state_masking = sca_lfsr_state;
  }
  if (context == kPentestLfsrOrder) {
    sca_lfsr_state_order = sca_lfsr_state;
  }
  return lfsr_out;
}

uint32_t pentest_linear_layer(uint32_t input) {
  uint32_t output =
      // output[7:0]
      (((input >> 0) & 0x1) << 7) | (((input >> 6) & 0x1) << 6) |
      (((input >> 11) & 0x1) << 5) | (((input >> 14) & 0x1) << 4) |
      (((input >> 1) & 0x1) << 3) | (((input >> 7) & 0x1) << 2) |
      (((input >> 10) & 0x1) << 1) | (((input >> 13) & 0x1) << 0) |
      // output[15:8]
      (((input >> 2) & 0x1) << 15) | (((input >> 4) & 0x1) << 14) |
      (((input >> 9) & 0x1) << 13) | (((input >> 12) & 0x1) << 12) |
      (((input >> 3) & 0x1) << 11) | (((input >> 5) & 0x1) << 10) |
      (((input >> 8) & 0x1) << 9) | (((input >> 15) & 0x1) << 8) |
      // output[23:16]
      (((input >> 16) & 0x1) << 23) | (((input >> 22) & 0x1) << 22) |
      (((input >> 27) & 0x1) << 21) | (((input >> 30) & 0x1) << 20) |
      (((input >> 17) & 0x1) << 19) | (((input >> 23) & 0x1) << 18) |
      (((input >> 26) & 0x1) << 17) | (((input >> 29) & 0x1) << 16) |
      // output[31:24]
      (((input >> 18) & 0x1) << 31) | (((input >> 20) & 0x1) << 30) |
      (((input >> 25) & 0x1) << 29) | (((input >> 28) & 0x1) << 28) |
      (((input >> 19) & 0x1) << 27) | (((input >> 21) & 0x1) << 26) |
      (((input >> 24) & 0x1) << 25) | (((input >> 31) & 0x1) << 24);

  return output;
}

uint32_t pentest_non_linear_layer(uint32_t input) {
  uint32_t output;
  if (input != 0) {
    // Perform the AES S-Box look ups bytewise.
    output = (uint32_t)(sbox[(input >> 24) & 0xFF] << 24) |
             (uint32_t)(sbox[(input >> 16) & 0xFF] << 16) |
             (uint32_t)(sbox[(input >> 8) & 0xFF] << 8) | sbox[input & 0xFF];
  } else {
    output = input;
  }

  return output;
}
